# 读取数据文件，转换成图片

import pickle
import glob
import os

import numpy as np
import cv2


# pickle是Python中的一种序列化格式，
# 它可以将Python对象转换为字节流并保存到文件中。
# 当需要使用这些对象时，可以从文件中读取字节流并重新构建出原始的Python对象。
def unpickle(file):
    with open(file, 'rb') as fo:
        my_dict = pickle.load(fo, encoding='bytes')  # 写的话是dump
    return my_dict


label_name = [
    'airplane',  # 飞机
    'automobile',  # 汽车
    'bird',  # 鸟
    'cat',  # 猫
    'deer',  # 鹿
    'dog',  # 狗
    'frog',  # 青蛙
    'horse',  # 马
    'ship',  # 船
    'truck'  # 卡车
]

# glob一个可以使用通配符查找文件与目录的库
test_list = glob.glob('../data/test_batch')
# print(train_list)

save_path = '../data/test'

for l in test_list:
    # print(l)
    l_dict = unpickle(l)
    # print(l_dict)
    # print(l_dict.keys())
    # dict_keys([b'batch_label', b'labels', b'data', b'filenames'])

    # enumerate可以同时列出数据和数据下标
    for im_idx, im_data in enumerate(l_dict[b'data']):
        # print(im_idx)
        # print(im_data)

        im_label = l_dict[b'labels'][im_idx]
        im_name = l_dict[b'filenames'][im_idx]

        # print(im_label,im_name,im_data)

        # 获得图片分类名字
        im_label_name = label_name[im_label]
        # 将图片变成 BGR
        im_data = np.reshape(im_data, (3, 32, 32))
        # 将图片变成 RGB
        im_data = np.transpose(im_data, (1, 2, 0))

        # 可视化
        # opencv读取的图片数值为BGR格式，与PIL库读取图片得到的RGB格式不同，因此如果用其他库的imshow函数来显示图片，图片的颜色会发生变化
        # cv2.imshow('im_data', cv2.resize(im_data, (200, 200)))
        # cv2.waitKey(0)

        # 存储
        if not os.path.exists('{}/{}'.format(save_path, im_label_name)):
            os.mkdir('{}/{}'.format(save_path, im_label_name))

        cv2.imwrite('{}/{}/{}'.format(
            save_path,
            im_label_name,
            im_name.decode('utf-8')),
            im_data)
